import os
import re
import requests
from urllib.parse import urljoin
from concurrent.futures import ThreadPoolExecutor
from Crypto.Cipher import AES  # 需要安装 pycryptodome
import ffmpy
import glob
import subprocess
import os
import tqdm
from requests.exceptions import HTTPError
# 安装依赖库：pip install pycryptodome requests
class download_m3u8():
    def safe_rename(old_path, new_name_gbk):
        """
        安全重命名文件，处理跨平台编码问题
        :param old_path: 原始文件路径（UTF-8 编码）
        :param new_name_gbk: 新文件名（需转换为 GBK 编码）
        """
        if not os.path.exists(old_path):
            raise FileNotFoundError(f"文件 {old_path} 不存在")
        
        # 根据系统处理新文件名编码
        new_name = new_name_gbk
        
        # 获取目录并生成新路径
        dir_name = os.path.dirname(old_path)
        new_path = os.path.join(dir_name, new_name)
        
        try:
            os.rename(old_path, new_path)
            print(f"重命名成功: {old_path} → {new_path}")
        except Exception as e:
            print(f"重命名失败: {str(e)}")
    def parse_m3u8(m3u8_url):
        """解析 m3u8 文件，提取 ts 地址、密钥 URL 和 IV"""
        try:
            response = requests.get(m3u8_url)
            response.raise_for_status()
            content = response.text

            # 提取密钥信息
            key_info = re.search(r'#EXT-X-KEY:METHOD=AES-128,URI="(.+?)",IV=(0x.+?)\s', content)
            key_uri = urljoin(m3u8_url, key_info.group(1)) if key_info else None
            iv_hex = key_info.group(2) if key_info else None

            # 转换 IV 为 bytes（移除 0x 前缀并按16字节对齐）
            iv = bytes.fromhex(iv_hex[2:]) if iv_hex else None

            # 提取所有 ts 文件的 URL（过滤注释行）
            lines = response.text.split('\n')
            # 使用正则匹配以 htpps 开头的行（不区分大小写）
            ts_urls = []
            for line in lines:
                #print(line,end='\n')
                if not line.startswith('#'):
                    ts_urls.append(line)
            #print(ts_urls)
            return ts_urls, key_uri, iv
        except HTTPError as e:
            print(f"HTTP 错误: {e}")
            print(f"响应内容: {response.text}")
        except requests.exceptions.RequestException as e:
            print(f"请求异常: {e}")
        return None, None, None
        

    def download_key(key_uri):
        """下载 AES 密钥"""
        response = requests.get(key_uri)
        response.raise_for_status()
        return response.content  # 密钥为 16 字节二进制数据

    def decrypt_ts(data, key, iv):
        """AES-128 CBC 模式解密"""
        cipher = AES.new(key, AES.MODE_CBC, iv=iv)
        decrypted_data = cipher.decrypt(data)
        return decrypted_data

    def download_ts(url, save_dir, index, total, key, iv):
        """下载并解密单个 ts 文件"""
        try:
            # 下载加密的 ts 数据
            response = requests.get(url, timeout=10)
            response.raise_for_status()
            encrypted_data = response.content

            # 解密数据
            decrypted_data = download_m3u8.decrypt_ts(encrypted_data, key, iv)

            # 保存解密后的 ts 文件
            filename = os.path.join(save_dir, f"decrypted_segment_{index:04d}.ts")
            with open(filename, 'wb') as f:
                f.write(decrypted_data)
            #print(f"已下载并解密 [{index}/{total}] {url}",end = '\r',flush = True)

        except Exception as e:
            print(f"失败 [{index}/{total}] {url}: {str(e)}")

    def download_all_ts(m3u8_url, save_dir='ts_files', max_workers=10):
        """主函数：处理加密视频下载"""
        os.makedirs(save_dir, exist_ok=True)

        # 解析 m3u8
        #print(m3u8_url)
        ts_urls, key_uri, iv = download_m3u8.parse_m3u8(m3u8_url)
        if not key_uri or not iv:
            print("m3u8 文件未包含有效的加密密钥或 IV")
            return

        # 下载密钥
        key = download_m3u8.download_key(key_uri)
        #print(f"密钥长度检查: {len(key)} 字节 (应为16字节)")

        # 多线程下载并解密
        total = len(ts_urls)
        #print(f"发现 {total} 个加密 ts 片段")

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            for index, url in enumerate(ts_urls, 1):
                executor.submit(download_m3u8.download_ts, url, save_dir, index, total, key, iv)

        #print("全部片段下载并解密完成！")
        return 1
    def merge_ts_files(save_dir, output_file="output.mp4"):
        """合并解密后的 ts 文件（需安装 FFmpeg）"""
        # 定义TS文件所在目录和输出MP4文件路径
        # 获取并按顺序排序.ts文件
        #ts_files = sorted(glob.glob('*.ts'), key=lambda x: int(x.split('_')[-1].split('.')[0]))
        current_directory = os.path.abspath(save_dir)

        # 获取所有TS文件的路径
        ts_files = [os.path.join(current_directory, f) for f in os.listdir(current_directory) if f.endswith('.ts')]
        ts_files.sort()
        file_list_txt = save_dir+'file_list.txt'
        #print(ts_files)
        # 生成FFmpeg所需的文件列表
        with open(file_list_txt, 'w') as f:
            for file in ts_files:
                f.write(f"file '{file}'\n")

        # 调用FFmpeg合并并输出为MP4
        subprocess.run([
            'ffmpeg',
            '-y',
            '-f', 'concat',      # 指定拼接格式
            '-safe', '0',        # 允许任意文件路径
            "-probesize", "10000000",       # 增加探测数据量
            "-analyzeduration", "10000000", # 延长分析时间
            '-i', file_list_txt, # 输入文件列表
            "-bsf:v", "h264_mp4toannexb",   # 修复 H.264 参数集
            '-c', 'copy',        # 直接复制流，无需重新编码
            '-f', 'mp4',  # 强制指定容器格式
            output_file         # 输出文件名
        ])

        os.remove(file_list_txt)
        for file_name in ts_files:
            if file_name.endswith('.ts'):
                os.remove(file_name)
